# FMCW雷达
import numpy as np
import matplotlib.pyplot as plt
from apps.fmcw.conf.app_const import AppConst as AC
from apps.fmcw.conf.app_config import AppConfig as AF

class FmcwRadar(object):
    

    @staticmethod
    def setup() -> None:
        for i in range(AF.numTX):
            AF.tx_loc[i][0], AF.tx_loc[i][1], AF.tx_loc[i][2] = i * AF.d_tx, 0, 0
        for i in range(AF.numRX):
            AF.rx_loc[i][0], AF.rx_loc[i][1], AF.rx_loc[i][2] = AF.tx_loc[AF.numTX-1][0] + AF.d_tx + i * AF.d_rx, 0, 0

    @staticmethod
    def draw_antennas() -> None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        for i in range(AF.numTX):
            ax.scatter(AF.tx_loc[i][0], AF.tx_loc[i][1], AF.tx_loc[i][2], color='blue', marker='o')
        for i in range(AF.numRX):
            ax.scatter(AF.rx_loc[i][0], AF.rx_loc[i][1], AF.rx_loc[i][2], color='red', marker='o')
        ax.set_xlabel('X-axis')
        ax.set_ylabel('Y-axis')
        ax.set_zlabel('Z-axis')
        ax.set_title('Transmitter and Receiver Locations')
        plt.show()



    # Complex signal
    def phase(tx, fx, slope):
        return 2 * np.pi * (fx * tx + slope / 2 * tx ** 2)

    def phase2(tx, fx, r, v):
        return 2 * np.pi * (2 * fx * r / AC.c + tx * (2 * fx * v / AC.c + 2 * AF.slope * r / AC.c))

    @staticmethod
    def step(times:np.ndarray, t_onePulse:np.ndarray) -> np.ndarray:
        AF.t_onePulse = t_onePulse
        for i in range(AF.numTX):
            for j in range(AF.numRX):
                AF.delays_tar1[i, j] = (np.linalg.norm(AF.tar1_loc - AF.rx_loc[j], ord=2, axis=1) + np.linalg.norm(AF.tar1_loc - AF.tx_loc[i], ord=2, axis=1))/AC.c
                AF.delays_tar2[i, j] = (np.linalg.norm(AF.tar2_loc - AF.rx_loc[j], ord=2, axis=1) + np.linalg.norm(AF.tar2_loc - AF.tx_loc[i], ord=2, axis=1))/AC.c
        fr1 = 2 * AF.r1[1] * AF.slope / AC.c
        fr2 = 2 * AF.r2[1] * AF.slope / AC.c
        fd1 = 2 * AF.v1_radial * AF.fc / AC.c  # doppler freq
        fd2 = 2 * AF.v2_radial * AF.fc / AC.c
        f_if1 = fr1 + fd1  # beat or IF freq
        f_if2 = fr2 + fd2
        for i in range(AF.numTX):
            for j in range(AF.numRX):
                for k in range(AF.numChirps):
                    phase_t = FmcwRadar.phase(t_onePulse, AF.fc, AF.slope)
                    phase_1 = FmcwRadar.phase(t_onePulse - AF.delays_tar1[i, j, k * AF.numADC:(k + 1) * AF.numADC], AF.fc, AF.slope)  # received
                    phase_2 = FmcwRadar.phase(t_onePulse - AF.delays_tar2[i, j, k * AF.numADC:(k + 1) * AF.numADC], AF.fc, AF.slope)
                    AF.signal_t[k * AF.numADC:(k + 1) * AF.numADC] = np.exp(1j * phase_t)
                    AF.signal_1[k * AF.numADC:(k + 1) * AF.numADC] = np.exp(1j * (phase_t - phase_1))
                    AF.signal_2[k * AF.numADC:(k + 1) * AF.numADC] = np.exp(1j * (phase_t - phase_2))
                AF.mixed[i, j] = AF.signal_1 + AF.signal_2
        # Initialize RDMs array
        mixed_t = AF.mixed.squeeze(0).transpose(1,0)
        RDC = None
        for i in range(AF.numADC):
            ddi = mixed_t[i::AF.numADC]
            if RDC is None:
                RDC = np.array([ddi])
            else:
                RDC = np.vstack((RDC, np.array([ddi])))
        return RDC
    
    @staticmethod
    def draw_signals() -> None:
        samples_num = 360
        # Plotting
        plt.figure()
        plt.subplot(3, 1, 1)
        plt.plot(AF.times, AF.signal_t.real)
        plt.title('TX')
        plt.xlim([AF.times[0], AF.times[0]+samples_num*AF.dt])
        plt.xlabel('Time (sec)')
        plt.ylabel('Amplitude')

        plt.subplot(3, 1, 2)
        plt.plot(AF.times, AF.signal_1.real)
        plt.title('RX')
        plt.xlim([AF.times[0], AF.times[0]+samples_num*AF.dt])
        plt.xlabel('Time (sec)')
        plt.ylabel('Amplitude')

        plt.subplot(3, 1, 3)
        plt.plot(AF.times, AF.mixed[0][7].real)
        plt.title('Mixed')
        plt.xlim([AF.times[0], AF.times[0]+samples_num*AF.dt])
        plt.xlabel('Time (sec)')
        plt.ylabel('Amplitude')

        plt.tight_layout()
        plt.show()
